#!/usr/bin/env python
# -*- coding: utf-8 -*-
from absl import app, flags
from datetime import datetime, date, time
from quant import strategy, utils
from rqalpha import run_func
from rqalpha.api import *
import pandas as pd
import pathlib
from matplotlib.dates import datestr2num
import json
from easydict import EasyDict
from typing import Dict
from ruamel import yaml
from ruamel.yaml.comments import CommentedSeq
import functools
import numpy as np

flags.DEFINE_string('name', None, 'Backtest config name')
flags.DEFINE_bool('eval', False, 'Eval and dump config')

FLAGS = flags.FLAGS


def runtest(bt: EasyDict):
    with pathlib.Path(__file__).parent.joinpath('rqcfg.yml').open(
            mode='r') as f:
        config = EasyDict(yaml.round_trip_load(f))
    config.base.start_date = bt.start_date
    config.base.end_date = bt.end_date
    config.base.frequency = bt.frequency
    config.mod_configs = {}

    vars_dict = EasyDict({})
    if bt.cache_dir is None:
        config.mod.local_data_source.enabled = False
    else:
        cache_dir = pathlib.Path(bt.cache_dir)
        if bt.instruments is None:
            instruments = None
        else:
            instruments = pd.read_csv(cache_dir.joinpath(bt.instruments),
                                      index_col=['order_book_id'])
            # Remove not listed, nan will cause rqalpha to fail
            instruments = instruments[~pd.isna(instruments.listed_date)]
            if 'maturity_date' in instruments.columns:
                instruments['maturity_day'] = instruments.maturity_date.apply(
                    lambda md: datestr2num(md))

        dfs = []
        for k, v in bt.bars.items():
            p = cache_dir.joinpath(k).expanduser()
            for f in p.parent.glob(p.name):
                df = pd.read_csv(f)
                df = df[df.datetime >= v.start]
                df = df[df.datetime <= v.end]
                dfs.append(df)
        bars = pd.concat(dfs)
        if bt.frequency == '1m':
            bars['datetime'] = pd.to_datetime(bars.datetime)
        else:
            bars['datetime'] = pd.to_datetime(bars.datetime)
            bars['datetime'] = bars.datetime.apply(
                lambda dt: datetime.combine(dt.date(), time(15, 0, 0)))
        if 'yield_to_maturity' in bars.columns:
            # e.g. 127013.XSHE
            bars = bars[~pd.isna(bars.yield_to_maturity)]
        bars.sort_values(['datetime', 'order_book_id'], inplace=True)
        bars.set_index(['datetime', 'order_book_id'], inplace=True)
        config.mod.local_data_source.bar_data = bars
        config.mod.local_data_source.instruments = instruments
        vars_dict.bars = bars
        vars_dict.instruments = instruments

    signals = {}
    results = {}
    if 'grid' in bt:
        vars_dict.grid = eval(bt.grid).tolist()
    for s in bt.strategies:
        for var in s.params:
            if var in vars_dict:
                s.params[var] = vars_dict[var]
        st = getattr(strategy, s.strategy)(**s.params)
        r = run_func(init=st.init,
                     handle_bar=st.handle_bar,
                     config=dict(config))
        results[st.name] = r['sys_analyser']
        if len(st.signals) > 0:
            if 'name' in s.params:
                signals[s.params.name] = pd.concat(st.signals)
            else:
                signals[s.strategy] = pd.concat(st.signals)

    if bt.plot:
        utils.plot_rqalpha_backtest_results(
            datetime.now(),
            bt.start_date,
            bt.end_date,
            results,
            savefile=pathlib.Path(__file__).parent.joinpath('%s.png' %
                                                            bt.name))
    return signals


def main(argv):
    del argv  # unused
    with pathlib.Path(__file__).parent.joinpath('backtest.yml').open(
            mode='r') as f:
        docs = {doc['name']: doc for doc in yaml.round_trip_load_all(f)}
    import logging
    if FLAGS.eval:
        for name in docs.keys() if FLAGS.name is None else [FLAGS.name]:
            bt = docs[name]
            if 'grid' in bt:
                grid = np.round(eval(bt['grid']), decimals=3)
                for s in bt['strategies']:
                    if 'grid' in s['params']:
                        s['params']['grid'] = CommentedSeq(grid.tolist())
                        s['params'].yaml_add_eol_comment(bt['grid'], 'grid')
        with pathlib.Path(__file__).parent.joinpath('backtest.yml').open(
                mode='w') as f:
            for bt in docs.values():
                f.write('---\n')
                f.write(yaml.round_trip_dump(bt, default_flow_style=True))
                f.write('\n')
    else:
        for name in docs.keys() if FLAGS.name is None else [FLAGS.name]:
            bt = docs[name]
            logging.info('Backtesting %s' % name)
            runtest(EasyDict(bt))


if __name__ == '__main__':
    app.run(main)
